package com.cdn.clock.component.rate;

import com.cdn.clock.component.IRateComponent;
import com.cdn.clock.config.LockConfig;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Collections;
import java.util.concurrent.TimeUnit;

/**
 * 蔡定努
 * 2024/07/06 11:25
 */

@ConditionalOnProperty(prefix = "clock", name = "rate-limit-type", havingValue = "redis", matchIfMissing = false)
public class RedisRateComponent implements IRateComponent {

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    private static final Log log = LogFactory.getLog(RedisRateComponent.class);


    public RedisRateComponent() {
        log.info("init RedisRateComponent success……………………");
    }

    private static final String LOCK_SCRIPT = "-- rate_limiter.lua\n" +
            "local windowKey = KEYS[1]\n" +
            "local window = tonumber(ARGV[1])\n" +
            "local count = tonumber(ARGV[2])\n" +
            "local currentCount = tonumber(redis.call(\"GET\", windowKey) or \"0\")\n" +
            "if currentCount < count then\n" +
            "    redis.call(\"INCRBY\", windowKey, 1)\n" +
            "    redis.call(\"EXPIRE\", windowKey, window)\n" +
            "    return 1 -- 允许访问\n" +
            "else\n" +
            "    return 0 -- 请求次数已满，拒绝访问\n" +
            "end";


    @Override
    public boolean rateLimit(String key, long count, long rateInterval, TimeUnit timeUnit) {
        String time = String.valueOf(timeUnit.toSeconds(rateInterval));
        String countStr = String.valueOf(count);

        RedisScript<Long> script = new DefaultRedisScript<>(LOCK_SCRIPT, Long.class);
        Long result = redisTemplate.execute(
                script,
                Collections.singletonList(key),
                time, // 时间窗口，单位秒
                countStr  // 时间窗口内允许的请求次数
        );
        return result != null && result == 1;
    }

    public boolean rateLimit(String key, int timeSecond, int count) {
        String time = String.valueOf(timeSecond);
        String countStr = String.valueOf(count);

        RedisScript<Long> script = new DefaultRedisScript<>(LOCK_SCRIPT, Long.class);
        Long result = redisTemplate.execute(
                script,
                Collections.singletonList(key),
                time, // 时间窗口，单位秒
                countStr  // 时间窗口内允许的请求次数
        );
        return result != null && result == 1;
    }




}
